import numpy as np
from skfuzzy import control as ctrl
from skfuzzy import membership

# 定义ANFIS结构
class ANFIS:
    def __init__(self):
        self.input1 = ctrl.Antecedent(np.arange(0, 10, 1), 'input1')
        self.input2 = ctrl.Antecedent(np.arange(0, 10, 1), 'input2')
        self.output = ctrl.Consequent(np.arange(0, 100, 1), 'output')

        self.input1.automf(3)
        self.input2.automf(3)
        self.output['low'] = membership.trimf(self.output.universe, [0, 0, 50])
        self.output['medium'] = membership.trimf(self.output.universe, [0, 50, 100])
        self.output['high'] = membership.trimf(self.output.universe, [50, 100, 100])

        rule1 = ctrl.Rule(self.input1['poor'] & self.input2['poor'], self.output['low'])
        rule2 = ctrl.Rule(self.input1['average'] & self.input2['average'], self.output['medium'])
        rule3 = ctrl.Rule(self.input1['good'] & self.input2['good'], self.output['high'])

        self.control_system = ctrl.ControlSystem([rule1, rule2, rule3])
        self.simulator = ctrl.ControlSystemSimulation(self.control_system)

    def evaluate(self, x1, x2):
        self.simulator.input['input1'] = x1
        self.simulator.input['input2'] = x2
        self.simulator.compute()
        return self.simulator.output['output']


# 定义卷积操作
def convolve_with_anfis(image, anfis_model, stride=1, padding=0):
    padded_image = np.pad(image, padding, mode='constant')
    output_shape = ((padded_image.shape[0] - 2) // stride + 1,
                    (padded_image.shape[1] - 2) // stride + 1)
    output_image = np.zeros(output_shape)

    for i in range(0, padded_image.shape[0] - 2 + 1, stride):
        for j in range(0, padded_image.shape[1] - 2 + 1, stride):
            region = padded_image[i:i + 3, j:j + 3]
            output_image[i // stride, j // stride] = anfis_model.evaluate(np.mean(region), np.std(region))

    return output_image


# 定义全连接层
def fully_connected_layer(input_data, weights, bias):
    return np.dot(input_data, weights) + bias


# 测试代码
if __name__ == '__main__':
    # 创建ANFIS模型
    anfis_model = ANFIS()

    # 测试图像
    test_image = np.random.rand(5, 5) * 10  # 5x5的随机图像

    # 第一层卷积
    conv1_output = convolve_with_anfis(test_image, anfis_model, stride=1, padding=1)
    print("第一层卷积后的图像：")
    print(conv1_output)

    # 第二层卷积
    conv2_output = convolve_with_anfis(conv1_output, anfis_model, stride=1, padding=1)
    print("第二层卷积后的图像：")
    print(conv2_output)

    # 展平输出以连接全连接层
    flattened_output = conv2_output.flatten()

    # 定义全连接层的权重和偏置
    weights = np.random.rand(flattened_output.size, 3)  # 假设输出维度为3
    bias = np.random.rand(3)

    # 全连接层输出
    fc_output = fully_connected_layer(flattened_output, weights, bias)
    print("全连接层输出：")
    print(fc_output)
